package com.xinxin.datastructure.tree;

/**
 * @author 史鑫鑫
 * @date Created in 2019/6/12 21:27
 */
@SuppressWarnings("unchecked")
public class SegmentTree2D<E> {
    /**
     * 节点内部类
     */
    private class Node {
        E e;
        int startX;
        int startY;
        int endX;
        int endY;
        Node child1;
        Node child2;
        Node child3;
        Node child4;

        Node(E e, int startX, int startY, int endX, int endY) {
            this.e = e;
            this.startX = startX;
            this.startY = startY;
            this.endX = endX;
            this.endY = endY;
            this.child1 = null;
            this.child2 = null;
            this.child3 = null;
            this.child4 = null;
        }

        Node(int startX, int startY, int endX, int endY) {
            this(null, startX, startY, endX, endY);
        }

        @Override
        public String toString() {
            return "Node{" +
                    "e=" + e +
                    ", startX=" + startX +
                    ", startY=" + startY +
                    ", endX=" + endX +
                    ", endY=" + endY +
                    '}';
        }
    }

    /**
     * 原数组
     */
    private E[][] data;
    /**
     * X长度
     */
    private int xLen;
    /**
     * Y长度
     */
    private int yLen;
    /**
     * 线段树
     */
    private Node root;
    /**
     * 合并规则
     */
    private Merger<E> merger;

    /**
     * 构造函数
     *
     * @param arr    二维数组
     * @param xLen   数组第一维长度
     * @param yLen   数组第二维长度
     * @param merger 合并规则
     */
    public SegmentTree2D(E[][] arr, int xLen, int yLen, Merger<E> merger) {
        if (xLen <= 0 || yLen <= 0 || arr == null) {
            throw new IllegalArgumentException("Array cannot be empty");
        }
        if (arr.length != xLen) {
            throw new IllegalArgumentException("xLen need equal arr's length");
        }
        this.merger = merger;
        this.xLen = xLen;
        this.yLen = yLen;
        this.data = (E[][]) new Object[xLen][yLen];
        for (int i = 0; i < xLen; i++) {
            if (arr[i].length != yLen) {
                throw new IllegalArgumentException("yLen need equal arr[i]'s length");
            }
            System.arraycopy(arr[i], 0, this.data[i], 0, yLen);
        }
        root = buildSegmentTree(0, 0, data.length - 1, data[0].length - 1);
    }

    /**
     * 构建(startX,startY)到(endX,endY)之间的线段树，并返回线段树的根节点
     *
     * @param startX 起点横坐标
     * @param startY 起点纵坐标
     * @param endX   终点横坐标
     * @param endY   终点纵坐标
     * @return 线段树根节点
     */
    private Node buildSegmentTree(int startX, int startY, int endX, int endY) {
        if (startX == endX && startY == endY) {
            return new Node(data[startX][startY], startX, startY, endX, endY);
        }
        int midX = startX + (endX - startX) / 2;
        int midY = startY + (endY - startY) / 2;
        Node node = new Node(startX, startY, endX, endY);
        node.child1 = buildSegmentTree(startX, startY, midX, midY);
        node.child2 = buildSegmentTree(midX + 1, startY, endX, midY);
        node.child3 = buildSegmentTree(startX, midY + 1, midX, endY);
        node.child4 = buildSegmentTree(midX + 1, midY + 1, endX, endY);
        node.e = merger.merge(merger.merge(node.child1.e, node.child2.e), merger.merge(node.child3.e, node.child4.e));
        return node;
    }

    /**
     * 返回区间内的元素个数
     *
     * @return 区间内的元素个数
     */
    public int getSize() {
        return xLen * yLen;
    }

    /**
     * 返回(xIndex,yIndex)位置的元素
     *
     * @param xIndex x索引
     * @param yIndex y索引
     * @return (xIndex, yIndex)位置的元素
     */
    public E get(int xIndex, int yIndex) {
        if (xIndex < 0 || xIndex >= xLen || yIndex < 0 || yIndex >= yLen) {
            throw new IllegalArgumentException("Index is illegal.");
        }
        return data[xIndex][yIndex];
    }

    /**
     * 返回区间[(queryStartX,queryStartY), (queryEndX,queryEndY)]的值
     *
     * @param queryStartX 开始横坐标
     * @param queryStartY 开始纵坐标
     * @param queryEndX   结束横坐标
     * @param queryEndY   结束纵坐标
     * @return 合并值
     */
    public E query(int queryStartX, int queryStartY, int queryEndX, int queryEndY) {
        // 参数不合法
        if (queryStartX < 0 || queryStartX >= xLen ||
                queryStartY < 0 || queryStartY >= yLen ||
                queryEndX < 0 || queryEndX >= xLen ||
                queryEndY < 0 || queryEndY >= yLen ||
                queryStartX > queryEndX || queryStartY > queryEndY) {
            throw new IllegalArgumentException("Index is illegal.");
        }
        return query(root, 0, 0, xLen - 1, yLen - 1, queryStartX, queryStartY, queryEndX, queryEndY);
    }

    /**
     * 在以node为根的线段树中[(startX, startY)...(endX, endY)]的范围里，搜索区间[(queryStartX, queryStartY)...(queryEndX, queryEndY)]的值
     *
     * @param node        根节点
     * @param startX      线段树开始横坐标
     * @param startY      线段树开始纵坐标
     * @param endX        线段树结束横坐标
     * @param endY        线段树结束纵坐标
     * @param queryStartX 搜索区间开始横坐标
     * @param queryStartY 搜索区间开始纵坐标
     * @param queryEndX   搜索区间结束横坐标
     * @param queryEndY   搜索区间结束纵坐标
     * @return 合并值
     */
    private E query(Node node, int startX, int startY, int endX, int endY, int queryStartX, int queryStartY, int queryEndX, int queryEndY) {
        if (startX == endX && startY == endY) {
            return node.e;
        }
        int midX = startX + (endX - startX) / 2;
        int midY = startY + (endY - startY) / 2;
        Node child1Node = node.child1;
        Node child2Node = node.child2;
        Node child3Node = node.child3;
        Node child4Node = node.child4;

        if (queryStartX >= midX + 1) {
            if (queryStartY >= midY + 1) {
                return query(child4Node, midX + 1, midY + 1, endX, endY, queryStartX, queryStartY, queryEndX, queryEndY);
            } else if (queryEndY <= midY) {
                return query(child2Node, midX + 1, startY, endX, midY, queryStartX, queryStartY, queryEndX, queryEndY);
            } else {
                E child2Result = query(child2Node, midX + 1, startY, endX, midY, queryStartX, queryStartY, queryEndX, midY);
                E child4Result = query(child4Node, midX + 1, midY + 1, endX, endY, queryStartX, midY + 1, queryEndX, queryEndY);
                return merger.merge(child2Result, child4Result);
            }
        } else if (queryEndX <= midX) {
            if (queryStartY >= midY + 1) {
                return query(child3Node, startX, midY + 1, midX, endY, queryStartX, queryStartY, queryEndX, queryEndY);
            } else if (queryEndY <= midY) {
                return query(child1Node, startX, startY, midX, midY, queryStartX, queryStartY, queryEndX, queryEndY);
            } else {
                E child1Result = query(child1Node, startX, startY, midX, midY, queryStartX, queryStartY, queryEndX, midY);
                E child3Result = query(child3Node, startX, midY + 1, midX, endY, queryStartX, midY + 1, queryEndX, queryEndY);
                return merger.merge(child1Result, child3Result);
            }
        } else {
            if (queryStartY >= midY + 1) {
                E child3Result = query(child3Node, startX, midY + 1, midX, endY, queryStartX, queryStartY, midX, queryEndY);
                E child4Result = query(child4Node, midX + 1, midY + 1, endX, endY, midX + 1, queryStartY, queryEndX, queryEndY);
                return merger.merge(child3Result, child4Result);
            } else if (queryEndY <= midY) {
                E child1Result = query(child1Node, startX, startY, midX, midY, queryStartX, queryStartY, midX, queryEndY);
                E child2Result = query(child2Node, midX + 1, startY, endX, midY, midX + 1, queryStartY, queryEndX, queryEndY);
                return merger.merge(child1Result, child2Result);
            } else {
                E child1Result = query(child1Node, startX, startY, midX, midY, queryStartX, queryStartY, midX, midY);
                E child2Result = query(child2Node, midX + 1, startY, endX, midY, midX + 1, queryStartY, queryEndX, midY);
                E child3Result = query(child3Node, startX, midY + 1, midX, endY, queryStartX, midY + 1, midX, queryEndY);
                E child4Result = query(child4Node, midX + 1, midY + 1, endX, endY, midX + 1, midY + 1, queryEndX, queryEndY);
                return merger.merge(merger.merge(child1Result, child2Result), merger.merge(child3Result, child4Result));
            }
        }
    }

    /**
     * 将(indexX, indexY)位置的值，更新为e
     *
     * @param indexX 行索引
     * @param indexY 列索引
     * @param e      更新值
     */
    public void set(int indexX, int indexY, E e) {
        if (indexX < 0 || indexX >= xLen || indexY < 0 || indexY >= yLen) {
            throw new IllegalArgumentException("Index is illegal");
        }
        data[indexX][indexY] = e;
        set(root, 0, 0, xLen - 1, yLen - 1, indexX, indexY, e);
    }

    /**
     * 在以node为根的线段树中更新index的值为e
     *
     * @param node   根节点
     * @param startX 区间开始行索引
     * @param startY 区间开始列索引
     * @param endX   区间结束行索引
     * @param endY   区间结束列索引
     * @param indexX 更新索引行坐标
     * @param indexY 更新索引列坐标
     * @param e      更新值
     */
    private void set(Node node, int startX, int startY, int endX, int endY, int indexX, int indexY, E e) {
        if (startX == endX && startY == endY) {
            node.e = e;
            return;
        }
        int midX = startX + (endX - startX) / 2;
        int midY = startY + (endY - startY) / 2;
        Node child1Node = node.child1;
        Node child2Node = node.child2;
        Node child3Node = node.child3;
        Node child4Node = node.child4;
        if (indexX >= midX + 1) {
            if (indexY >= midY + 1) {
                set(child4Node, midX + 1, midY + 1, endX, endY, indexX, indexY, e);
            } else {
                set(child2Node, midX + 1, startY, endX, midY, indexX, indexY, e);
            }
        } else {
            if (indexY >= midY + 1) {
                set(child3Node, startX, midY + 1, midX, endY, indexX, indexY, e);
            } else {
                set(child1Node, startX, startY, midX, midY, indexX, indexY, e);
            }
        }
        node.e = merger.merge(merger.merge(node.child1.e, node.child2.e), merger.merge(node.child3.e, node.child4.e));
    }


    private String toString(Node node) {
        if (node == null) {
            return "null ";
        }
        return node + " " + toString(node.child1) + toString(node.child2) + toString(node.child3) + toString(node.child4);
    }

    @Override
    public String toString() {
        return toString(root);
    }
}
